#!/usr/bin/python
# multio3.py - align omega output on all rigid ring fragments
# jji jan 04
# based on templates from OpenEye SS
#adpoted => can read multimol2 file, changes solvation table if multiple molecule entries were generated
#problem: if ring is not flat, mol2db gets confused => have to be seperated (independend, if ring is used as anchor or if ring is attached to anchor)
#an issue: mol2db only allows label with 9 digits -> where to add my mutiple_ring_identifier?
#=> had to sacrify first two digits, add ring identifier at end
#=> use fix_db_labels.py to remove delete ring_identifier label

#make sure that the first atom in the mol2 file is part of the rigid fragment, otherwise mol2db might get confused


#looses molecules that don't have a ring!

#can not handle duplicates in file names <------------------------------

import os, sys, chemistry
from openeye.oechem import *
#ifs_single = oemolistream(sys.argv[1])

rot=DoubleArray(9)
trans=DoubleArray(3)

if not (len(sys.argv) == 5 or len(sys.argv)== 6):
	print 'multio3_mod_mult.py  <directory with omega files> <directory with amsol output directories (*CAV)> <only planar rings True|False (meaning less for pattern matching)> <zinc titles True|False> <file with smarts> (only for pattern matching)'
	sys.exit(1)

omega_files = os.listdir(sys.argv[1])
amsol_dir = sys.argv[2]

if len(sys.argv) == 6:
	pattern_matching = True
	smarts_file = open(sys.argv[5], 'r')
	rms_cut_off = 0.001 # fewer atoms in pattern matching -> must be more accurate, otherwise mol2db screws up
	print "pattern matching"
else:
	pattern_matching = False
	print "ring matching"
	only_planar_rings = sys.argv[3]
	rms_cut_off = 0.01
	if only_planar_rings == 'True':
		only_planar_rings = True
	else:
		only_planar_rings = False

#zinc titles have to be treated diffrently, because they are already "pre-adjusted" (they are not unique when you download them
zinc_title = sys.argv[4]
if zinc_title == 'True':
	zinc_title = True
	print "processing zinc files"
else:
	zinc_title = False
	print "processing non-zinc files"




mol = OEGraphMol()
mol2 = OEGraphMol()
solv_dict = {}


#tmp = oemolostream('tmp.mol2')
#tmp2 = oemolostream('tmp2.mol2')
tmp_mol_list = []
tmp2_mol_list = []


#---------------------
def read_ensemble(ifs_mult,tmp_mol_list, tmp2_mol_list):
	mol = OEGraphMol()

	#tmp2.close()
	#tmp2 = oemolistream('tmp2.mol2')
	#OEReadMolecule(tmp2, mol)
	#tmp2.close()
	mol = tmp2_mol_list[0]
	title = mol.GetTitle()


	#reset all files
	#tmp.close()
	#command = 'rm -f tmp.mol2'
	#os.system(command)
	#tmp = oemolostream('tmp.mol2')
	tmp_mol_list = []
	#tmp2 = oemolostream('tmp2.mol2')
	tmp2_mol_list = []
	#OEWriteMolecule(tmp,mol)
	tmp_mol_list.append(mol)
	
	
	new_title = title	
	while new_title == title and ifs_mult.IsValid():
		mol = OEGraphMol() #if this statement is not inserted the molecules in the list will be all be identical
		OEReadMolecule(ifs_mult, mol)
		new_title = mol.GetTitle()
		#mol.SetTitle(str(counter))
		#print title, new_title
		if new_title == title:
			#OEWriteMolecule(tmp, mol)
			tmp_mol_list.append(mol)
		else:
			#OEWriteMolecule(tmp2, mol)
			tmp2_mol_list.append(mol)
	return tmp_mol_list, tmp2_mol_list
	

#---------------------
def gen_new_title(title, counter, zinc_label):
	if zinc_label:
		#it is an already adjusted title
		old_number = int(title[4]) * 10 #this assumes that there are a maximum of nine versions of the same molecules in Zinc
		new_number = old_number + counter 
		new_title = 'MFC' + str(new_number).rjust(2) + title[5:]
	else:
		new_title = 'MFC' + title[5:] + str(counter).rjust(2)

	new_title = new_title.replace(' ', '0')
	return new_title

#-------------------------------
def read_solv_table(solv_table_file):
	dict = {}
	for i in solv_table_file.xreadlines():
		if i[0:1] <> ' ':
			end = i.find(' ')
			identifier = i[:end]
			dict[identifier] = [i,[]]
		else:
			dict[identifier][1].append(i)
	
	return dict
		

		
	
#---------------------
def find_atoms_in_plane(ringlist,mol, count):
	#for now, I'll only extend planar ring systems
	pred = OEPartPredAtom(ringlist)
  	refmol = OECreateOEMol()
 	
	old_list = ringlist
	new_list = ringlist

	
	for i in xrange(1,count+1):
		new_list = []
		checked_atoms = []
		checked_neighbours = []
		for j in old_list:
			new_list.append(j)
			start = True
   		pred.SelectPart(i)
    		OESubsetMol(refmol, mol, pred) 
		if chemistry.check_planar(refmol):
			while (new_list <> old_list) or start:
				#print 'new loop'
				old_list = []
				start = False
				for j in new_list:
					old_list.append(j)
				for atom in mol.GetAtoms():
					if old_list[atom.GetIdx()] == i and atom.GetIdx() not in checked_atoms:
						checked_atoms.append(atom.GetIdx())
						#atom is part of the current ring_system
						for neighbours in atom.GetAtoms() :
							#print 'neighbours'
							if old_list[neighbours.GetIdx()] == 0 and neighbours.GetIdx() not in checked_neighbours:
							#atom must not be part of the current ringsystem or another ring_system
								checked_neighbours.append(neighbours.GetIdx())
								#test if this atom is in the same plane
								tmp_list = []
								for j in new_list:
									tmp_list.append(j)
								tmp_list[neighbours.GetIdx()] = i
								pred2 = OEPartPredAtom(tmp_list) 
								pred2.SelectPart(i)
								temp_mol = OECreateOEMol()
								OESubsetMol(temp_mol, mol, pred2)
								if chemistry.check_planar(temp_mol):
									#keep this atom
									#print 'planar'
									new_list = []
									for j in tmp_list:
										new_list.append(j)
	return new_list

					
	
#---------------------
def super_impose_only_on_planar_rings(ringlist,mol, count):
	#for now, I'll only extend planar ring systems
	pred = OEPartPredAtom(ringlist)
  	refmol = OECreateOEMol()
 	
	old_list = ringlist
	ring_count = 0

	new_list = []
	for j in old_list:
		new_list.append(j)
	
	for i in xrange(1,count+1):
   		pred.SelectPart(i)
    		OESubsetMol(refmol, mol, pred) 
		if chemistry.check_planar(refmol):
			#print "keep this ring"
			#print ringlist
			ring_count = ring_count + 1
			
		else:
			#print "reject this ring"
			#delete this ring from ring_list
			for atom_number in range(0,len(new_list)):
				#print new_list[atom_number]
				if new_list[atom_number] == i:
					new_list[atom_number] = 0
					#print 'test'
				#adjust ring_numbers
				elif new_list[atom_number] > i:
					new_list[atom_number] = new_list[atom_number] -1
					
			

	#print new_list
	return new_list, ring_count

					
	
#---------------------
def OrderMol2db(mol2, pred):
	index_list = []
	atom_list_match = []
	for atom in mol2.GetAtoms(pred):
	#print atom.GetIdx()
		atom_list_match.append(atom)
		index_list.append(atom.GetIdx())
	for atom in mol2.GetAtoms():
		if atom not in atom_list_match:
			atom_list_match.append(atom)
			index_list.append(atom.GetIdx())

	mol2.OrderAtoms(atom_list_match)
	return mol2, index_list

#---------------------
def  get_matches(mol, patlist):
	#generate empty list
	pattern_matches = []

	match = 0
	for pat,smarts,description in patlist:
		for matchbase in pat.Match(mol,True) :
			#print description
			#generate empty list
			pattern_match = []
			for index in range(1, mol.NumAtoms() + 1):
				pattern_match.append(0)

			match = match + 1
			#print match
			for matchpair in matchbase.GetAtoms():
				#print matchpair.target.GetIdx(), matchpair.target.GetName()
				pattern_match[matchpair.target.GetIdx()] = match
			pattern_matches.append(pattern_match)
      


	
	#for i in pattern_matches:
		#print i

	return pattern_matches, match
#---------------------

#MAIN

if pattern_matching:
	#some options from the original match.py file
	usa=1; asym=1; exph=0; max=0; smartsfile=None; verbose=0;
	qmolfile=None; kekule=0; atom_expr=None; bond_expr=None;
	smartslist=None; qmollist=None; nowarn=0;

	smarts_list=chemistry.get_smartslist_file(smarts_file)

	patlist= chemistry.get_patlist(smarts_list, atom_expr, bond_expr,max)



for file in omega_files:
	if file [-8:] == 'mol2.bz2':
		command = 'bunzip2 ' + sys.argv[1] + '/' + file
		os.system(command)
		file = file[:-4]
 	if file[-4:] == 'mol2':
		file_name = sys.argv[1] + '/' + file
		ifs_mult = oemolistream(file_name)
		out_file = file_name[:-5] + '_mult_rings.mol2'
		ofs_mult = oemolostream(out_file)
		solv_table_file_name = amsol_dir + '/' + file[0:-7] + 'os_amsol_cav_CYC/' + file[0:-7] +'os_CYC.solv'
		#print solv_table_file_name
		new_solv_table_file_name = amsol_dir + '/' + file[0:-7] + 'os_amsol_cav_CYC/' + file[0:-7]  +'mult_rings_os_CYC.solv'
		solv_table_file = open(solv_table_file_name, 'r')
		new_solv_talbe_file = open(new_solv_table_file_name, 'w')


		#read solvation_table, this might have to be changed, if tables get too long to be held in memory
		solv_dict = read_solv_table(solv_table_file)
		#print solv_dict.keys()


		#Start: read first molecule and write it to tmp2 file
		OEReadMolecule(ifs_mult, mol)
		#OEWriteMolecule(tmp2, mol)
		tmp2_mol_list=[mol]
		go_on = True
		count = 0  #needed for empty files

		while go_on:
			tmp_mol_list, tmp2_mol_list = read_ensemble(ifs_mult,tmp_mol_list, tmp2_mol_list)
			#read first molecule of ensemble in tmp, determine superposition stuff, write out
			#tmp =oemolistream('tmp.mol2')
			#OEReadMolecule(tmp, mol)
			#tmp.close()
			mol = tmp_mol_list[0]

 	 		title =  mol.GetTitle()
			#print mol.GetTitle()
			#go only on, if solvation energy calculation was sucessfull
			if solv_dict.has_key(title):
				if not pattern_matching:
 		 			count, ringlist = OEDetermineRingSystems(mol)
					if only_planar_rings:
						new_list, count = super_impose_only_on_planar_rings(ringlist,mol, count)
						#print count
					else:
						new_list = find_atoms_in_plane(ringlist,mol,count)
					match_list = []
					for i in range(1,count+1):
						match_list.append(new_list)
				else:
					#print title
					match_list, count = get_matches(mol, patlist)
					#print match_list

  				refmol = OECreateOEMol()

				super_pos_counter = 0
				#print count, ' count'
				for i in range(1,count+1):
					pred = OEPartPredAtom(match_list[i-1])
    					pred.SelectPart(i)
    					OESubsetMol(refmol, mol, pred) 
					ref_mol_list = []	
					ref_mol_list.append(refmol)
    					#sys.stderr.write('align on fragment #%d ' % i)
    					#ifs2 = oemolistream('tmp.mol2')
  				        #find out, how many reference molecules are needed (because some rings are not planar)
					#this has also to be done for pattern matching because not all patterns are polar					
					if not pattern_matching:
						planar = chemistry.check_planar(refmol)
					else:
						planar = False # always check if pattern matching
					if not planar:
					   #print mol.GetTitle()
      					   #while OEReadMolecule(ifs2, mol2):
      					   for mol2 in tmp_mol_list:
      						thismol = OECreateOEMol()
      						OESubsetMol (thismol, mol2, pred)
						#print 'Check reference molecules'
						found = False
						for refmol in ref_mol_list:
							if not found:
      								rms=OERMSD(refmol,thismol,0,1,1,rot,trans)
								#print 'rmsd %5.2f\n' % rms
								if rms <= rms_cut_off:
									found = True
									#print found
						#further reference molecules needed?
						if not found:
							refmol = OECreateOEMol()
							OESubsetMol(refmol, mol2, pred)
							ref_mol_list.append(refmol)

					for refmol in ref_mol_list:
					   super_pos_counter = super_pos_counter + 1
					   #print super_pos_counter
					  
 					   #ifs2.close()
					   #ifs2 = oemolistream('tmp.mol2')
					   #while OEReadMolecule(ifs2, mol2):
 					   for mol2 in tmp_mol_list:
						#print 'here', mol2.GetTitle()
     						thismol = OECreateOEMol()
     						OESubsetMol (thismol, mol2, pred)
      						rms=OERMSD(refmol,thismol,0,1,1,rot,trans)
						#print rms
						if (rms <= rms_cut_off): 
      							#print 'rmsd %5.2f\n' % rms
      							OERotate(mol2,rot)
      							OETranslate(mol2,trans)

							#for zinc, you have to adjust the titles before => don't destroy this adjustment here
							new_title = gen_new_title(title, super_pos_counter, zinc_title)

							mol2.SetTitle(new_title)

							#order atoms for mol2db
							mol2, index_list = OrderMol2db(mol2, pred)						
 		 					
      							OEWriteMolecule(ofs_mult, mol2)
							#print mol2.GetTitle()
					  #identifier must be unique
					   #new_solv_talbe_file.write('MFC' + solv_dict[title][0][4:12] + str(super_pos_counter) + solv_dict[title][0][12:])	
					   #print index_list
					   new_solv_talbe_file.write(new_title + solv_dict[title][0][12:])
					   for index in index_list:
					   	new_solv_talbe_file.write(solv_dict[title][1][index])
				



			if not ifs_mult.IsValid():
				go_on = False
				#Here I would lose the last molecule if it is rigid

				#ifs = oemolistream('tmp2.mol2')
				last_mol = OEGraphMol()
				#OEReadMolecule(ifs, last_mol)
				#if the last molecule is an ensemble => tmp2_list is empty
				if len(tmp2_mol_list) > 0:
					last_mol = tmp2_mol_list[0]
					title = last_mol.GetTitle()
					new_title = gen_new_title(title, 1, zinc_title)
					last_mol.SetTitle(new_title)	

					#order atoms for mol2db
					last_mol, index_list = OrderMol2db(last_mol, pred)				
					OEWriteMolecule(ofs_mult, last_mol)
					new_solv_talbe_file.write(new_title + solv_dict[title][0][12:])	
					for index in index_list:
				   		new_solv_talbe_file.write(solv_dict[title][1][index])

				
				

		new_solv_talbe_file.close()
		solv_table_file.close()
		command ='nice bzip2 ' + file_name
		#os.system(command)
		command = 'nice bzip2 ' +  file_name[:-5] + '_mult_rings.mol2'
		#os.system(command)

